import numpy as np
import pandas as pd
from sklearn.preprocessing import StandardScaler
from sklearn.impute import SimpleImputer
from sklearn.preprocessing import LabelEncoder  # 修正：LabelEncoder属于preprocessing模块
from sklearn.metrics import classification_report
from sklearn.model_selection import train_test_split
from sklearn.neighbors import KNeighborsClassifier
from sklearn.svm import SVC
from sklearn.linear_model import LogisticRegression  # 新增：导入逻辑回归
import joblib


class Classification:
    def __init__(self):
        self.df = pd.read_csv('fina_indicator.csv')

    def get_conditions(self):
        df = self.df
        
        numeric_features = ['eps', 'total_revenue_ps', 'undist_profit_ps', 'gross_margin', 
                           'fcff', 'fcfe', 'tangible_asset', 'bps', 'grossprofit_margin', 'npta',
                           'max_close', 'min_close', 'the_close']  
        
        # 初始化填充器
        imputer = SimpleImputer(strategy='mean')
        df[numeric_features] = imputer.fit_transform(df[numeric_features])
        
        # 计算比率特征
        df['max_ratio'] = df['max_close'] / df['the_close']
        df['min_ratio'] = df['min_close'] / df['the_close']
        
        # 确定阈值
        high_return_threshold = df['max_ratio'].quantile(0.4)
        high_risk_threshold = df['min_ratio'].quantile(0.4)
        
        # 分类条件
        conditions = [
            (df['max_ratio'] >= high_return_threshold) & (df['min_ratio'] <= high_risk_threshold),
            (df['max_ratio'] < high_return_threshold) & (df['min_ratio'] >= high_risk_threshold),
            (df['max_ratio'] >= high_return_threshold) & (df['min_ratio'] > high_risk_threshold),
            (df['max_ratio'] < high_return_threshold) & (df['min_ratio'] < high_risk_threshold),
        ]
        choices = ['高风险低收益', '低风险高收益', '高风险高收益', '低风险低收益']
        df['category'] = pd.Categorical(np.select(conditions, choices, default='其他'))
        
        # 特征与标签处理
        features = ['eps', 'total_revenue_ps', 'undist_profit_ps', 'gross_margin', 
                   'fcff', 'fcfe', 'tangible_asset', 'bps', 'grossprofit_margin', 'npta']
        label_encoder = LabelEncoder()
        df['category_encoded'] = label_encoder.fit_transform(df['category'])
        
        # 特征标准化
        scaler = StandardScaler()
        X = scaler.fit_transform(df[features])
        Y = df['category_encoded']
        
        # 划分数据集
        X_train, X_test, Y_train, Y_test = train_test_split(
            X, Y, test_size=0.2, random_state=42
        )
        
        # 保存预处理工具
        joblib.dump(scaler, "scaler.joblib")
        joblib.dump(imputer, "imputer.joblib")
        return X_train, X_test, Y_train, Y_test, label_encoder
    
    def logistic_regression_utils(self, X_train, X_test, Y_train, Y_test, label_encoder):
        log_reg = LogisticRegression(max_iter=1000, random_state=42)
        log_reg.fit(X_train, Y_train)
        
        # 预测与评估
        y_pred = log_reg.predict(X_test)
        print("逻辑回归分类报告：")
        print(classification_report(Y_test, y_pred, target_names=label_encoder.classes_))
        
        # 保存模型
        joblib.dump(log_reg, 'logistic_regression_model.joblib')
        joblib.dump(label_encoder, 'label_encoder.joblib') 

    def knn_utils(self, X_train, X_test, Y_train, Y_test, label_encoder):
        knn = KNeighborsClassifier(n_neighbors=5)
        knn.fit(X_train, Y_train)
        
        y_pred = knn.predict(X_test)
        print("KNN 分类报告：")
        print(classification_report(Y_test, y_pred, target_names=label_encoder.classes_))
        
        joblib.dump(knn, 'knn_model.joblib')
        joblib.dump(label_encoder, 'label_encoder.joblib')

    def svc_utils(self, X_train, X_test, Y_train, Y_test, label_encoder):
        svc = SVC(kernel='linear')
        svc.fit(X_train, Y_train)
        
        y_pred = svc.predict(X_test)
        print("SVC 分类报告：")
        print(classification_report(Y_test, y_pred, target_names=label_encoder.classes_))
        
        joblib.dump(svc, 'svc_model.joblib')
        joblib.dump(label_encoder, 'label_encoder.joblib')


if __name__ == '__main__':
    c1 = Classification()
    X_train, X_test, Y_train, Y_test, label_encoder = c1.get_conditions()
    c1.svc_utils(X_train, X_test, Y_train, Y_test, label_encoder)
    c1.knn_utils(X_train, X_test, Y_train, Y_test, label_encoder)
    c1.logistic_regression_utils(X_train, X_test, Y_train, Y_test, label_encoder) 